#!/usr/bin/python
# -*- encoding: utf-8 -*-

"""
@File:      tc_ai_tensorflow2_distribute.py
@Time:      2023/07/4 11:12:50
@Author:    zhangtaibo
@Version:   1.0
@Contact:   zhangtaibo.ztb@alibaba-inc.com
@License:   Mulan PSL v2
@Modify:    zhangtaibo
"""
from common.basetest import LocalTest

class Test(LocalTest):
    """
    See tc_ai_tensorflow2_distribute.yaml for details

    :avocado: tags=fix,P1,noarch,local,tensorflow2_ai_container
    """
    PARAM_DIC = {"pkg_name": "tensorflow2"}
    def setUp(self):
        super().setUp()
        ret_c,self.docker_name = self.cmd(self.container_engine+" ps |grep %s |awk '{print $NF}'" % self.version)
        self.cmd(self.container_engine+" cp res/ai/tensorflow/tensorflow2_distribute.py %s:/tmp" % self.docker_name)

    def test(self):
        _, output = self.cmd("python3 /tmp/tensorflow2_distribute.py",container_flag=1)
        self.assertTrue("loss" in output and "accuracy" in output, 'tensorflow2 distribute test failed, please check.')

    def tearDown(self):
        super().tearDown()